import argparse
import textwrap
import shutil
import binascii
import os
import hashlib
import struct


def get_args():
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
                                     description=textwrap.dedent('''
                                     A tool to recover raw img'''))
    # input for soc
    parser.add_argument('-img', required=True, dest='img',
                        help='INPUT: img with header, Parse N/E from it')
    parser.add_argument('-raw', required=False, dest='raw',
                        help='INPUT: recovered raw img')
    parser.add_argument('--hash', help="cal the hash value N|E, and transfer the efuse burned words",
                        action="store_true")
    parser.add_argument('--rcvr', help="recover the raw img without header",
                        action="store_true")
    parser.add_argument('--cmp', help="cmp raw img and hash in img with header",
                        action="store_true")
    return parser.parse_args()


def __write_raw_img(raw, img, code_len):
    raw.seek(0)
    img.seek(0x4000)
    rsv_len = code_len
    while rsv_len > 0:
        if rsv_len > 4096:
            raw.write(img.read(4096))
        else:
            raw.write(img.read(rsv_len))
        rsv_len -= 4096


def __cal_hash(buf_n, buf_e):
    sha256_hash = hashlib.sha256()
    sha256_hash.update(buf_n)
    sha256_hash.update(buf_e)
    hash_buf = sha256_hash.digest()
    print(binascii.hexlify(hash_buf))
    return hash_buf[::-1]


def __cal_image_hash(f):
    sha256_hash = hashlib.sha256()
    for byte_block in iter(lambda: f.read(4096), b""):
        sha256_hash.update(byte_block)
    f.seek(0)
    return binascii.hexlify(sha256_hash.digest())


def __shift_hash(hash_buf):
    rotpk = []
    shift_rotpk = [0] * 9
    for i in xrange(0, 32, 4):
        rotpk.append(struct.unpack('<I', hash_buf[i : i + 4])[0])
    print(map(hex, rotpk))
    tmp = 0x0
    for j in xrange(0, 8):
        shift_rotpk[j] = ((rotpk[j] & 0x1) << 31) | tmp
        tmp = (rotpk[j] >> 1)
        print("the {0}-th word is {1}".format(j, hex(shift_rotpk[j])))
    shift_rotpk[8] = tmp
    print("the 8-th word is " + hex(shift_rotpk[8]))


def main():
    args = get_args()
    if args.rcvr:
        with open(args.img, 'rb') as img:
            img.seek(0x460)
            code_len = struct.unpack('<I', img.read(4))[0]
            tmp_file = args.raw + '.tmp'
            with open(tmp_file, 'wb+') as raw:
                __write_raw_img(raw, img, code_len)
        shutil.copyfile(tmp_file, args.raw)
        if os.path.exists(tmp_file):
            os.remove(tmp_file)

    if args.hash:
        with open(args.img, 'rb') as r_f:
            r_f.seek(0x5c)
            buf_e = r_f.read(512)
            r_f.seek(0x25c)
            buf_n = r_f.read(512)
            hash_buf = __cal_hash(buf_n, buf_e)
            __shift_hash(hash_buf)

    if args.cmp:
        with open(args.img, 'rb') as r_f:
            r_f.seek(0x2c)
            org_hash = binascii.hexlify(r_f.read(32))
            with open(args.raw, 'rb') as raw:
                if cmp(org_hash, __cal_image_hash(raw)):
                    print('hash cmp fail, recover fail')
                else:
                    print('hash cmp succeed, recover succeed')


if __name__ == '__main__':
    main()


